package com.shujia.spark.mllib

import org.apache.spark.ml.classification.{LogisticRegression, LogisticRegressionModel}
import org.apache.spark.sql.{DataFrame, Dataset, Row, SparkSession}

object Demo4TrainMode {
  def main(args: Array[String]): Unit = {
    val spark: SparkSession = SparkSession
      .builder()
      .master("local[8]")
      .appName("image")
      .getOrCreate()
    import spark.implicits._

    //1、读取图片数据
    val dataDF: DataFrame = spark.read
      .format("libsvm")
      .option("numFeatures", 784)
      .load("data/image_data")

    //2、将数据切分成训练集和测试集
    //训练集负责训练模型，测试集负责测试模型的准确率
    val Array(train, test) = dataDF.randomSplit(Array(0.8, 0.2))

    //3、选择算法
    val regression = new LogisticRegression()

    //4、将训练集带入算法训练模型
    val model: LogisticRegressionModel = regression.fit(train)

    //5、将测试集带入模型评估模型准确率
    val testDF: DataFrame = model.transform(test)

    testDF.show(false)

    //6、计算准确率
    val p: Double = testDF.where($"label" === $"prediction").count().toDouble / testDF.count() * 100
    println(s"准确率：$p")

    //7、保存模型
    model.save("data/image_model")
  }
}
